/**
 * WebSocket service for AI content generation streaming.
 */

import { WebSocketService } from './websocket.service';

export interface GenerationMessage {
  type: 'token' | 'error' | 'complete' | 'cancelled' | 'start' | 'heartbeat' | 'heartbeat_ack';
  content?: string;
  error?: string;
  timestamp: string;
  metadata?: {
    model?: string;
    tokens_used?: number;
    message_id?: string;
    token_count?: number;
    char_count?: number;
    duration?: number;
  };
}

export interface GenerationRequest {
  template_id: string;
  variables: Record<string, any>;
  model?: string;
}

export class GenerationWebSocketService extends WebSocketService {
  private taskId: string | null = null;
  private isGenerating = false;
  private generationHandlers: {
    onToken?: (token: string) => void;
    onComplete?: (metadata: any) => void;
    onError?: (error: string) => void;
    onStart?: (metadata: any) => void;
  } = {};

  constructor(taskId: string, token?: string) {
    const wsUrl = import.meta.env.VITE_WS_URL || 'ws://localhost:8002';
    const tokenParam = token ? `?token=${encodeURIComponent(token)}` : '';
    
    super({
      url: `${wsUrl}/api/v1/streaming/ws/generation/${taskId}${tokenParam}`,
      reconnectInterval: 2000,
      maxReconnectAttempts: 5,
      heartbeatInterval: 30000,
      maxReconnectDelay: 10000,
      reconnectBackoffMultiplier: 1.5,
    });

    this.taskId = taskId;
    this.setupMessageHandlers();
  }

  /**
   * Start content generation.
   */
  async startGeneration(request: GenerationRequest): Promise<void> {
    if (!this.isConnected()) {
      await this.connect();
    }

    this.isGenerating = true;
    this.send('generate', {
      ...request,
      task_id: this.taskId,
    });
  }

  /**
   * Cancel ongoing generation.
   */
  cancelGeneration(): void {
    if (this.isGenerating) {
      this.send('cancel', { task_id: this.taskId });
      this.isGenerating = false;
    }
  }

  /**
   * Set generation event handlers.
   */
  setGenerationHandlers(handlers: {
    onToken?: (token: string) => void;
    onComplete?: (metadata: any) => void;
    onError?: (error: string) => void;
    onStart?: (metadata: any) => void;
  }): void {
    this.generationHandlers = handlers;
  }

  /**
   * Setup internal message handlers.
   */
  private setupMessageHandlers(): void {
    // Handle different message types
    this.on('message', (message: GenerationMessage) => {
      switch (message.type) {
        case 'token':
          if (message.content && this.generationHandlers.onToken) {
            this.generationHandlers.onToken(message.content);
          }
          break;

        case 'complete':
          this.isGenerating = false;
          if (this.generationHandlers.onComplete) {
            this.generationHandlers.onComplete(message.metadata);
          }
          break;

        case 'error':
          this.isGenerating = false;
          if (message.error && this.generationHandlers.onError) {
            this.generationHandlers.onError(message.error);
          }
          break;

        case 'start':
          if (this.generationHandlers.onStart) {
            this.generationHandlers.onStart(message.metadata);
          }
          break;

        case 'cancelled':
          this.isGenerating = false;
          if (this.generationHandlers.onComplete) {
            this.generationHandlers.onComplete({ cancelled: true });
          }
          break;

        case 'heartbeat':
          // Respond to heartbeat
          this.send('heartbeat', {});
          break;
      }
    });
  }

  /**
   * Get current generation status.
   */
  isCurrentlyGenerating(): boolean {
    return this.isGenerating;
  }

  /**
   * Get task ID.
   */
  getTaskId(): string | null {
    return this.taskId;
  }
}

/**
 * Factory function to create generation WebSocket service.
 */
export function createGenerationWebSocket(
  taskId: string,
  token?: string
): GenerationWebSocketService {
  return new GenerationWebSocketService(taskId, token);
}